#!/usr/bin/env python
# -*- coding:utf-8 -*-

from __future__ import print_function

import timeit

import numpy

import theano
import theano.tensor as T
from theano.tensor.shared_randomstreams import RandomStreams

from load_dataset import *
import logging
from test_logging import logging_conf
from kmeans_alg import *
from sklearn import metrics
from scipy.cluster.vq import *

logging_conf()
logger = logging.getLogger('SGAE.py')

class GAE(object):

    def __init__(
            self,
            numpy_rng,
            theano_rng=None,
            input=None,
            n_visible=195,
            n_hidden=128,
            W=None,
            bhid=None,
            bvis=None
    ):

        self.n_visible = n_visible
        self.n_hidden = n_hidden

        # create a Theano random generator that gives symbolic random values
        if not theano_rng:
            theano_rng = RandomStreams(numpy_rng.randint(2 ** 30))

        # note : W' was written as `W_prime` and b' as `b_prime`
        if not W:
            initial_W = numpy.asarray(
                numpy_rng.uniform(
                    low=-4 * numpy.sqrt(6. / (n_hidden + n_visible)),
                    high=4 * numpy.sqrt(6. / (n_hidden + n_visible)),
                    size=(n_visible, n_hidden)
                ),
                dtype=theano.config.floatX
            )
            W = theano.shared(value=initial_W, name='W', borrow=True)

        if not bvis:
            bvis = theano.shared(
                value=numpy.zeros(
                    n_visible,
                    dtype=theano.config.floatX
                ),
                borrow=True
            )

        if not bhid:
            bhid = theano.shared(
                value=numpy.zeros(
                    n_hidden,
                    dtype=theano.config.floatX
                ),
                name='b',
                borrow=True
            )

        self.W = W
        self.b = bhid
        self.b_prime = bvis
        self.W_prime = self.W.T
        self.theano_rng = theano_rng

        # if no input is given, generate a variable representing the input
        if input is None:
            self.x = T.dmatrix(name='input')
        else:
            self.x = input

        self.params = [self.W, self.b, self.b_prime]

    def get_hidden_values(self, input):
        """ Computes the values of the hidden layer """
        return T.nnet.sigmoid(T.dot(input, self.W) + self.b)

    def get_reconstructed_input(self, hidden):
        """Computes the reconstructed input given the values of the
        hidden layer

        """
        return T.nnet.sigmoid(T.dot(hidden, self.W_prime) + self.b_prime)

    def get_cost_updates(self, learning_rate, semi_data):
        """ This function computes the cost and the updates for one trainng
        step of the dA """

        self.y = self.get_hidden_values(self.x)
        z = self.get_reconstructed_input(self.y)

        # object function: cost entropy
        # L = - T.sum(self.x * T.log(z) + (1 - self.x) * T.log(1 - z), axis=1)
        # cost = T.mean(L)

        # object function: euclidean distance   + 1000 * (self.y.T * semi_data * self.y).trace()
        a = T.sqrt(T.sum((self.x - z) ** 2))
        b = T.dot(T.dot(self.y.T, semi_data), self.y).trace()
        cost = a + 1000 * b

        # compute the gradients of the cost of the `dA` with respect
        # to its parameters
        gparams = T.grad(cost, self.params)
        # generate the list of updates
        updates = [
            (param, param - learning_rate * gparam)
            for param, gparam in zip(self.params, gparams)
            ]

        return (cost, updates, a, b)


def test_ae(learning_rate=0.1, iter_num=10000, error_threshold=0.000001, community_num=5, dataset_name='cornell', test_num=1, borrow=True):

    # print parameters log information
    logger.info("Start...")
    logger.info("there are param's value, learning_rate:%f, iter_num:%d, error_threshold:%f"
                % (learning_rate, iter_num, error_threshold))

    # load data set
    datasets = load_cornell("../data/cornell/cornell_adj.txt", "../data/cornell/cornell_content.txt", dataset_name)

    for test_index in range(test_num):
        # define
        row_num = len(datasets)
        node_num = row_num/2
        index = T.lscalar()
        x = T.matrix('x')
        data = theano.shared(numpy.asarray(datasets, dtype=theano.config.floatX), borrow=borrow)
        in_layer_size = [195, 128]
        out_layer_size = [128, 64]
        n_layer = len(out_layer_size)
        params_list = []
        every_layer_output_list = []

        # start train model
        start_time = timeit.default_timer()
        for i in range(n_layer):

            rng = numpy.random.RandomState()
            theano_rng = RandomStreams(rng.randint(2 ** 30))  # random integer [0, 2 ** 30)

            da = GAE(
                numpy_rng=rng,
                theano_rng=theano_rng,
                input=x,
                n_visible=in_layer_size[i],
                n_hidden=out_layer_size[i]
            )

            temp_data = load_constraint(node_num)
            cost, updates, a, b = da.get_cost_updates(
                learning_rate=learning_rate,
                semi_data=temp_data
            )

            train_da = theano.function(
                [index],
                [cost, a, b],
                updates=updates,
                givens={
                    x: data[index: (index + row_num)]
                }
            )

            stopping = False
            epoch = 1
            while not stopping:
                # go through trainng set
                c = []
                # for batch_index in range(sample_num):
                # print ('start train batch %d' % batch_index)
                batch_index = 0
                cost_val, a_val, b_val = train_da(batch_index)
                c.append(cost_val)

                cost_result = numpy.mean(c)
                print('Training layer %i, epoch %d, cost %f, a %f, b %f'
                      % ((i + 1), epoch, cost_result, numpy.mean(a_val), numpy.mean(b_val)))

                epoch += 1
                if epoch > iter_num:
                    stopping = True

            params_list.append(da.params)
            y = da.get_hidden_values(data)
            every_layer_output_list.append(y.eval())
            data.set_value(y.eval())

        end_time = timeit.default_timer()

        training_time = (end_time - start_time)
        logger.info(('The train code runtime is %.2fm' % ((training_time) / 60.)))

        # the last output value
        last_output_result = every_layer_output_list[-1]

        # the result is (2n * t), we need to sum the first (n * t) and second (n * t)
        # because the structure and the content are similar
        print(len(last_output_result))
        cluster_input_data = []
        for t_index in range(node_num):
            cluster_input_data.append(last_output_result[t_index] + last_output_result[t_index + node_num])

        # k-means algorithm
        kmeans_centroid, kmeans_lables = kmeans2(data=numpy.array(cluster_input_data),
                                                 k=community_num, iter=20, minit='points')
        print("cluster result: ")
        logging.info(kmeans_lables + 1)
        cornell_ground_truth = load_cornell_groundtruth("../data/cornell/cornell_ground_truth.txt")
        logging.info(metrics.normalized_mutual_info_score(cornell_ground_truth, kmeans_lables))


if __name__ == '__main__':
    test_ae()
